import sys
# sys.path.append(r'/home/user10000630/notespace/Chatbot_227_ef9a')
sys.path.append('..')
import pandas as pd
from site_packages.utils.job import DataOp
import os
from sklearn.model_selection import train_test_split
from configs.settings import DATA_PATH_BERT


def read_files():
    files = ['toy_train.txt', 'toy_valid.txt']

    dfs = []
    for file in files:
        df = pd.read_csv(os.path.join(DATA_PATH_BERT, file), sep='\t', header=None, names=['question', 'answer'])
        dfs.append(df)
        
    return pd.concat(dfs)


def split_data(df, train_size=0.7, valid_size=0.2):
    assert train_size > 0 and valid_size > 0 and (train_size + valid_size) < 1, "比例传入错误"
    df = df.sample(frac=1).reset_index(drop=True)
    row_cnt = df.shape[0]
    train_upper_index = int(row_cnt * train_size)
    valid_upper_index = train_upper_index + int(row_cnt * valid_size)
    train = df.loc[:train_upper_index, :].copy()
    valid = df.loc[train_upper_index:valid_upper_index, :].copy()
    test = df.loc[valid_upper_index:, :].copy()
    DataOp.save(train, 'bert_train',is_model=False)
    DataOp.save(valid, 'bert_valid',is_model=False)
    DataOp.save(test, 'bert_test',is_model=False)
    
    
if __name__ == '__main__':
    df = read_files()
    split_data(df)